import requests
import json
from typing import Iterable, List

#HOST="http://quickstart-20250205-zfc6.1587392929778144.cn-wulanchabu.pai-eas.aliyuncs.com/"         # 替换为实际 AI 服务地址
#SK="YTc2MmI0NDRiODE4ZjJlNGMxYmIzYjUwNTdmMmU2YWUyN2M5NDNmNg=="   # 替换为实际的授权密钥
HOST="http://quickstart-20250206-oone.1587392929778144.cn-wulanchabu.pai-eas.aliyuncs.com/"
SK="Y2FjMjc5NjczNjIwOGM3N2E2OTNhNjAwMWFlOGZmYzhhMmZiMGM4Ng=="
class AIClient:
    def __init__(self, 
                 system_prompt: str = "你是树轮科技部署的deepseek大模型."):
        self.host = HOST
        self.secret_key = SK
        self.system_prompt = system_prompt

    def post_http_request(self, prompt: str,
                          history: List[tuple] = None,
                          max_new_tokens: int = 20480,
                          temperature: float = 0.8,
                          top_k: int = 1,
                          top_p: float = 0.8,
                          langchain: bool = False,
                          use_stream_chat: bool = False) -> requests.Response:
        if history is None:
            history = []

        headers = {
            "User-Agent": "Test Client",
            "Authorization": f"{self.secret_key}"
        }

        pload = {
            "prompt": prompt,
            "system_prompt": self.system_prompt,
            "top_k": top_k,
            "top_p": top_p,
            "temperature": temperature,
            "max_new_tokens": max_new_tokens,
            "use_stream_chat": use_stream_chat,
            "history": history
        }
        if langchain:
            pload["langchain"] = langchain

        response = requests.post(self.host, headers=headers, json=pload, stream=use_stream_chat)
        return response
    '''
    def get_streaming_response(self, response: requests.Response) -> Iterable[str]:
        """
        This function yields streamed responses from AI API as the response is being generated.
        """
        for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False, delimiter=b"\0"):
            if chunk:
                try:
                    data = json.loads(chunk.decode("utf-8"))
                    output = data["response"]
                    history = data["history"]
                    print(output)
                    yield output, history
                except json.JSONDecodeError as e:
                    print(f"Error decoding JSON: {e}")
                    continue
    '''
    def get_streaming_response(self, response: requests.Response) -> Iterable[tuple]:
        previous_output = ""
        for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False, delimiter=b"\0"):
            if chunk:
                try:
                    data = json.loads(chunk.decode("utf-8"))
                    # 从服务器返回的数据中获取累积的文本
                    output = data["response"]
                    history = data["history"]
                    # 如果新的累积文本 output 不以之前的 previous_output 开头，
                    # 说明服务器已经重置了累积文本（例如，新的一轮对话开始了），
                    # 那么将 previous_output 重置为空字符串。
                    if not output.startswith(previous_output):
                        previous_output = ""
                    # 计算新增的部分：当前返回的 output 去除之前已经输出的部分
                    incremental = output[len(previous_output):]
                    # 更新保存的 previous_output 为当前的完整文本
                    previous_output = output
                    print(incremental)  # 调试时打印新增加的部分
                    yield incremental, history
                except json.JSONDecodeError as e:
                    print(f"Error decoding JSON: {e}")
                    continue

